from XEdu.hub import Workflow as wf
import numpy as np
import cv2
#from playsound import playsound
import pygame
import threading
import time

flag = 0
text = ' '
num = 1
sound_no = 'paimen-no.wav'
sound_ok = 'paimen-ok.wav'
pygame.mixer.init()

def play_sound(num,sound_no,sound_ok):
    if num == 0:
        pygame.mixer.music.load(sound_no)
    elif num == 1:
        pygame.mixer.music.load(sound_ok)
    pygame.mixer.music.play()
    while pygame.mixer.music.get_busy():
        pygame.time.Clock().tick(10)  # 限制循环的频率，避免CPU过载
    time.sleep(5)

body = wf(task='body26' )
det_body = wf(task='det_body')
basenn = wf(task='basenn',checkpoint='model/20240522.onnx')
cap = cv2.VideoCapture(0)
while True:
    ret,img = cap.read()
    if ret:
        height, width, channels = img.shape#获取图片的尺寸，为防止图片分辨率的不同，影响判断的结果
        det,det_box = det_body.inference(data=img,img_type='cv2') # 进行模型推理
        pose,pose_img = body.inference(data=img,img_type='cv2')
        if det != []:
            result = body.inference(data=img)
            result = np.concatenate(result).reshape(1, -1)#将推理结果重整为一行的一维数组
            result[::2]= result[::2]/width#数组中的单数值除以宽度
            result[1::2]= result[1::2]/height#数组中的双数值除以高度

            res = basenn.inference(data=result)
            res_show = basenn.format_output(lang="en")
            label = ['tuobei','ok','']
            if res_show[0]['scores'] > 0.85:
                num = res_show[0]["prediction"]#rea_show是一个二维数组，其中第一的值为一个字典
        
            position = (50, 100) # 左上角起始点的坐标
            font = cv2.FONT_HERSHEY_SIMPLEX
            scale = 1
            color = (255, 255, 255) # BGR颜色值（白色）
            thickness = 2
 
            # 将文字写入图像中
            if num == num:#手势不变，计数加一
                flag = flag+1
            else:
                flag = 0#否则清零
            if flag >=30:#大于30，输出
                print(num)
                flag = 0
                text = label[num]#文本内容

                
                
                sound_thread = threading.Thread(target=play_sound, args=(num,sound_no,sound_ok,))
                sound_thread.start()
                
            cv2.putText(pose_img, text, position, font, scale, color, thickness)

        
            # 显示结果图像
            cv2.imshow("Image with Text", pose_img)
            if cv2.waitKey(1) & 0xFF == 27:
                break
cap.release()
cv2.destroyAllWindows